/*
 * Wazuh Vulnerability Scanner - Unit Tests
 * Copyright (C) 2015, Wazuh Inc.
 * May 2, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#include "remediationDataCache_test.hpp"
#include "MockSocketDBWrapper.hpp"
#include "TrampolineSocketDBWrapper.hpp"

/**
 * @brief Compare two remediation objects
 *
 * @param remediation1
 * @param remediation2
 * @return true if the remediations are equal
 * @return false if the remediations are not equal
 */
bool remediationsAreEqual(const Remediation& remediation1, const Remediation& remediation2)
{
    if (remediation1.hotfixes.size() != remediation2.hotfixes.size())
    {
        return false;
    }

    for (const auto& hotfix : remediation1.hotfixes)
    {
        if (remediation2.hotfixes.find(hotfix) == remediation2.hotfixes.end())
        {
            return false;
        }
    }

    return true;
}

TEST_F(RemediationDataCacheTest, InsertMultipleItems)
{
    RemediationDataCache<TrampolineSocketDBWrapper> cache;
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();
    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_)).Times(0);

    std::string agentId {"1"};

    {
        // Set value in cache
        Remediation remediationData {
            .hotfixes = {"hotfix1", "hotfix2"},
        };

        cache.addRemediationData(agentId, remediationData);

        // Get value from cache
        const auto retrievedData = cache.getRemediationData(agentId);

        // Verify that the returned value is the same as the one set
        EXPECT_TRUE(remediationsAreEqual(remediationData, remediationData));
    }

    {
        // Set another value in cache
        Remediation remediationData {
            .hotfixes = {"hotfix3", "hotfix4"},
        };

        cache.addRemediationData(agentId, remediationData);

        // Get value from cache
        const auto retrievedData = cache.getRemediationData(agentId);

        // Verify that the returned value is equal to both the values set
        Remediation expected {.hotfixes = {"hotfix1", "hotfix2", "hotfix3", "hotfix4"}};
        EXPECT_TRUE(remediationsAreEqual(retrievedData, expected));
    }
}

TEST_F(RemediationDataCacheTest, SetAndGetSuccess)
{
    RemediationDataCache<TrampolineSocketDBWrapper> cache;
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_)).Times(2);

    std::string agentId {"1"};

    // Try to get value from empty cache
    EXPECT_TRUE(remediationsAreEqual(cache.getRemediationData(agentId), Remediation {}));

    // Set value in cache
    Remediation remediationData {
        .hotfixes = {"hotfix1", "hotfix2"},
    };

    cache.addRemediationData(agentId, remediationData);

    // Verify that the returned value is the same as the one set
    EXPECT_TRUE(remediationsAreEqual(cache.getRemediationData(agentId), remediationData));

    // Try to get from non existing agent
    EXPECT_TRUE(remediationsAreEqual(cache.getRemediationData("2"), Remediation {}));
}

TEST_F(RemediationDataCacheTest, EmptyResponseFromDB)
{
    RemediationDataCache<TrampolineSocketDBWrapper> cache;
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    nlohmann::json queryResponse = nlohmann::json::array();

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(1)
        .WillOnce(testing::SetArgReferee<1>(queryResponse));

    std::string agentId {"1"};

    EXPECT_TRUE(remediationsAreEqual(cache.getRemediationData(agentId), Remediation {}));
}

TEST_F(RemediationDataCacheTest, ResponseFromDB)
{
    RemediationDataCache<TrampolineSocketDBWrapper> cache;
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    nlohmann::json queryResponse = R"([{"hotfix": "hotfix1"}, {"hotfix": "hotfix2"}])"_json;

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(1)
        .WillOnce(testing::SetArgReferee<1>(queryResponse));

    std::string agentId {"1"};

    // Get value from the cache
    const auto remediationData = cache.getRemediationData(agentId);

    // Verify that the returned value is the same as the one set
    Remediation expected {.hotfixes = {"hotfix1", "hotfix2"}};

    EXPECT_TRUE(remediationsAreEqual(remediationData, expected));
}

TEST_F(RemediationDataCacheTest, UnrecoverableExceptionOnDB)
{
    RemediationDataCache<TrampolineSocketDBWrapper> cache;
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(1)
        .WillOnce(testing::Throw(std::runtime_error("Error on DB")));

    std::string agentId {"1"};

    // Attempt to get value from the cache
    EXPECT_THROW(cache.getRemediationData(agentId), std::runtime_error);
    spSocketDBWrapperMock.reset();
}

TEST_F(RemediationDataCacheTest, RecoverableExceptionOnDB)
{
    RemediationDataCache<TrampolineSocketDBWrapper> cache;
    spSocketDBWrapperMock = std::make_shared<MockSocketDBWrapper>();

    EXPECT_CALL(*spSocketDBWrapperMock, query(testing::_, testing::_))
        .Times(1)
        .WillOnce(testing::Throw(SocketDbWrapperException("Warning on DB")));

    std::string agentId {"1"};

    // Attempt to get value from the cache
    EXPECT_THROW(cache.getRemediationData(agentId), WdbDataException);
    spSocketDBWrapperMock.reset();
}
